package edu.gmu.atelier;

/**
 * Performs LU decomposition on the passed matrix.  Computes the
 * upper triangular matrix U by forward substitution, along the way
 * the lower triangular matrix L is updated with the inverse eliminations.
 * <p>
 * Given Ax=b, solving for x is done in two steps.  During forward 
 * elimination the equation is transformed Ax=b -> U'x'=y.  Note that
 * the b does not affect what L and U' will look like and we can use
 * the same LU' to solve for different b's.
 * <p>
 * We use partial pivoting by columns, therefore the computed U' = UP and
 * the computed x' = P^Tx such that LU'x'=b.
 * <p>
 * First step is to compute y.  Considering Ax=b LU'x'=b Ly=b, we know
 * L and b and can solve for y using forward substitution.
 * <p>
 * Second step is to compute x', we know U' and now y, solve for x' in U'x'=y
 * using back substitution.
 * <p>
 * Finally the computed x' = P^Tx has the be multiplied by P to return x.
 * 
 * @author James H. Pope
 */
public class LUFactor implements Solver
{
    private Matrix a  = null;
    private Matrix up = null;
    private Matrix l  = null;
    
    // I'd like to make this a simple array or better yet set of ei vectors
    private Matrix p  = null;
    
    public LUFactor( Matrix a )
    {
        this.a = a;
        this.computeLU();
    }
    
    //------------------------------------------------------------------------//
    // Interface methods
    //------------------------------------------------------------------------//
    public Matrix getP()
    {
        return this.p;
    }
    
    public Matrix getUP()
    {
        return up;
    }
    
    public Matrix getU()
    {
        // Note: U = UP P^T
        return up.mult(p.transpose());
    }
    
    public Matrix getL()
    {
        return this.l;
    }
    
    /**
     * Computed x1 may be different from x if round-off errors
     * and is likely to occur.  Assume the error was introduced as z1
     * therefore x1 = x - z1.  We know that Ax1 = b1 -> A(x-z1) = b1
     * -> b - b1 = Az1 -> r1 = Az1 where r1 = b - b1, the residual.
     * Solve for z1 where Az1 = r1, then the "actual" x = x1+z1.
     * Of course this new "x" from z1 also has errors and we iterate.
     * <p>
     * Another simpler derivation is r1 = b - Ax1 = Az1 where z1 = x - x1.
     * <p>
     * Note that the improvement uses the already existing LU decomposition
     * for A in each iteration for solving Az=r along with Axi=bi.
     * <p>
     * The method returns when one of three conditions is met
     *   1. maxIterations have been executed
     *   2. ri is small enough according to e
     *   3. ri starts to diverge 
     * 
     * @param x
     * @param b
     * @param maxIter
     * @param e
     * @return xi such that Axi closer to b than Ax
     */
    public Vector improve( Vector x, Vector b, int maxIter, Epsilon e )
    {
        //Epsilon e = Epsilon.E; //Epsilon.makeBinary(53);
        //int MAX_ITER = 3;
        Vector xi = x;
        double riPrevNorm = 0.0;
        //System.out.println("Improve x1="+x1);
        for (int i = 0; i < maxIter; i++)
        {
            Vector bi = a.mult(xi);
            Vector ri = b.sub(bi); // dont get these swapped! will move away
            double riCurrNorm = ri.norm2().doubleValue();
            if(i==0)
            {
                // prime previous for first time around
                riPrevNorm = riCurrNorm;
            }
            // If r1 is similar 0, then x similar xi so return
            // We use double machine precision
            //System.out.println("Iter="+i);
            if( e.similar(riCurrNorm, 0.0) == true )
            {
                //System.out.println("    b's are close enough bi="+bi+", b="+b);
                break;
            }
            // If we are moving away (i.e. delta getting larger) bail
            // We return the lowest value before diverging, however, that
            // does not mean we found the xi closest to x. Should >= ???
            if( riCurrNorm > riPrevNorm )
            {
                //System.out.println("    diverged riCurrNorm="+riCurrNorm+", riPrevNorm="+riPrevNorm);
                break;
            }
            riPrevNorm = riCurrNorm;
            Vector zi = this.solve(ri);
            xi = xi.add(zi);
        }
        return xi;
    }
    
    /**
     * Solves the equation Ax=LU'x'=b.  Since U'x'=y, Ly=b, so we first solve for
     * y from L and specified b.  Then we solve for x' in equation U'x'=y.
     * Then finally return x = P x' = P P^Tx.
     * @return x
     */
    public Vector solve( Vector b )
    {
        Vector y  = LUFactor.forwardSubstitution(l, b);
        Vector x1= LUFactor.backSubstitution(up, y);
        
        /*
         * Took me a while to figure out why this is correct.
         * 
         * Typical row permuation does left multiply and only affects b
         *      Ax = c1a1 + c2a2 + c3a3 =  b, where b = {b1,b2,b3}
         *     PAx = c1a1 + c2a2 + c3a3 = Pb, where b = {b2,b1,b3}
         * In order to "undo" to get un-permuted b, left multiply b by P^T
         *     P^T Pb = b 
         * 
         * We used col permuation which only affects x
         *      LUIx = c1a1 + c2a2 + c3a3 =  b, where    x = {c1,c2,c3}
         *   LUPP^Tx = c2a2 + c1a1 + c3a3 =  b, where P^Tx = {c2,c1,c3}
         * In order to "undo" to get the un-permuted x, left multiply x by P
         *     P P^Tx = x 
         * 
         * The actual "u" computed was UP.
         * The actual "x" computed was P^Tx
         * 
         * The equation Ly=b is not affected.
         * The equation Ux=y is affected
         * 
         *     U'   x' = y
         *    UP P^Tx  = y
         *       P^Tx  = (UP)^-1 y
         *       P^Tx  = P^T U^-1 y
         * 
         * Now we know why this works. Typically when using row permuations
         * we undo the permuation on b using P^T before solving.
         * We used col permuations so we undo the permutation on x using P
         * after solving.
         */
        return x1.mult(p);
    }
    
    /**
     * Gets the determinant of the matrix A (actually determined by using
     * the already factored upper triangular U and simply multiplying the
     * diagonal).
     * @return det(A)
     */ 
    public double determinant()
    {
        double det;

        // We have already checked that matrix A is square so this is safe.
        int maxSize = this.up.getRows();
        det = this.up.getVector(0).getElement(0).doubleValue();
        for(int i = 1; i<maxSize; i++)
        {
            det = det * this.up.getVector(i).getElement(i).doubleValue();
        }

        return det;
    }
    
    private void computeLU( )
    {   
        //--------------------------------------------------------------------//
        // Make u - starting out a copy of a
        //--------------------------------------------------------------------//
        // Make local reference to u,l,p and set at end
        int n = a.getCols();
        int m = a.getRows();
        Matrix up= a.copy();
        Matrix p = a.getFactory().makeIdentity(n);
//        java.util.List<Vector> es = new java.util.ArrayList<Vector>();
//        for (int i = 0; i < n; i++)
//        {
//            es.add( new StandardVector(i,n) );    
//        }
//        Matrix p = new Matrix(es, a.getFactory());
        Matrix l = a.getFactory().makeIdentity(n);
        
        for (int c = 0; c < n-1; c++)
        {            
            //----------------------------------------------------------------//
            // Partial Pivoting (by col), find "largest" pivot value in row
            // Actually two ideas here
            // 1. If pivot zero, swap with column with a non-zero pivot
            // 2. Always use largest pivot to avoid round off errors
            // Largest means in magnitude, so absolute value is used.
            // Performing the second also performs the first.
            //----------------------------------------------------------------//
            int s = c;
            Element largestValue = up.getValue(c, c).abs();
            for (int z = c; z < n; z++)
            {
                Element possiblePivot = up.getValue(c, z).abs();
                if( possiblePivot.compareTo( largestValue ) > 0 )
                {
                    largestValue = possiblePivot;
                    s   = z;
                }
            }
            
            if(c!=s)
            {
                //System.out.println("Swapping cols "+c+ " and "+s);
                up.swapCol(c, s);
                //p.swap(i, s);
                
                /*
                 * Need to resolve, rather use permuation matrix but will
                 * need to solve x.dot(p.transpose), but can't currently
                 * cannot right multiply a vector by a matrix.
                 */
//                // Update the permutation "matrix"
//                int ival = this.perms[c];
//                this.perms[c] = this.perms[s];
//                this.perms[s] = ival;
                
                // Matrix version
                p.swapCol(c, s);
            }
            //----------------------------------------------------------------//
            
            // Forward elimination
            Element pivotCoeff = up.getValue(c, c);
            // If the pivot is zero, the Matrix is singular
            if( pivotCoeff.isZero() )
            {
                String msg = "The matrix is singular and cannot be LU "+
                             "factored, pivot zero at "+c;
                //System.out.println(up);
                throw new IllegalArgumentException(msg);
            }
            
            for (int r = c+1; r < n; r++)
            {
                Element lowerCoeff = up.getValue(r, c);
                Element multiplier = pivotCoeff.inverse().mult(lowerCoeff);
                up.mulAddTo( c, multiplier.negate(), r );
                l.setValue(r, c, multiplier);
            }
        }
        
        // Now set on instance, can be referenced later
        this.up= up;
        this.l = l;
        this.p = p;
        //System.out.println("U:"+up);
        //System.out.println("L:"+l);
        //System.out.println("P:"+p);
        
//        System.out.println("Perms:");
//        for (int i = 0; i < perms.length; i++)
//        {
//            System.out.println("      "+i+"="+perms[i]);
//        }
//        System.out.println("\n");
    }
    
    /**
     * Takes matrix u in upper triangle form and column vector y of constants
     * and solves the system of equations by back-substitution placing
     * the solution in the returned colVector x.
     * <p>
     * Or more simply put, solves for x in the equation Ux=y.
     * @param u matrix assumed to be in upper triangle form
     * @param y constants vectors
     * @return x solution
     */
    public static Vector backSubstitution( Matrix u, Vector y )
    {
        int n = y.getLength();
        Vector x = y.copy();
        Matrix identity = u.copy();

        for (int i = n-1; i >= 0; i--)
        {
            Element pivotCoefficient = identity.getValue(i, i);

            // Divide both u and b by pivotCoefficient (copies in our case)
            Element inversePivot = pivotCoefficient.inverse();
            identity.mulTo( i, inversePivot );
            x.mulTo(        i, inversePivot);

            for( int r = i-1; r >= 0; r-- )
            {
                // Determine the multiplier by pivot coefficient
                Element upperCoefficient = identity.getValue(r, i);
                //Element upperCoefficient = identity.getValue(r, i);
                // We know that the pivotCoefficient is now one (though
                // may not be due to round-off, hack to be one???)
                //double multiplier = upperCoefficient / pivotCoefficient;
                Element multiplier = upperCoefficient.negate();

                identity.mulAddTo( i, multiplier, r);
                x.mulAddTo(        i, multiplier, r);
            }
        }
        //System.out.println( "Identity: " + identity );

        return x;
    }

    /**
     * Takes matrix l in lower triangle form and column vector b of constants
     * and solves for y by forward-substitution.
     * <p>
     * Or more simply put, solves for y in the equation Ly=b.
     * @param l
     * @param b
     * @return y
     */
    public static Vector forwardSubstitution( Matrix l, Vector b )
    {
        int n = b.getLength();
        Vector y = b.copy();
        Matrix identity = l.copy();

        for (int i = 0; i < n-1; i++)
        {
            // We know that all the pivot coefficients are one since L
            // Just to be save, div first and then mulAdd
            Element pivotCoefficient = identity.getValue(i, i);

            identity.mulTo( i, pivotCoefficient);
            y.mulTo(        i, pivotCoefficient);

            for (int r = i+1; r < n; r++)
            {
                // Determine the multiplier by pivot coefficient
                Element lowerCoefficient = identity.getValue(r, i);
                // We know that the pivotCoefficient is now one (though
                // may not be due to round-off, hack to be one???)
                //double multiplier = upperCoefficient / pivotCoefficient;
                Element multiplier = lowerCoefficient.negate();

                identity.mulAddTo( i, multiplier, r);
                y.mulAddTo(        i, multiplier, r);
            }
        }
        //System.out.println( "Identity: " + identity );

        return y;
    }
    
    
    public String toString()
    {
        StringBuilder buf = new StringBuilder();
        buf.append( "Matrix a: " );
        buf.append( a.toString() );
        buf.append( "Matrix u: " );
        buf.append( up.toString() );
        buf.append( "Matrix l: " );
        buf.append( l.toString() );
        return buf.toString();
    }
    
}
